/**
 * 版权所有(C)，上海海鼎信息工程股份有限公司，2025，所有权利保留。
 * <p>
 * 项目名： cdx-ai-agent
 * 文件名： RerAdvisor.java
 * 模块说明：
 * 修改历史：
 * 2025年09月23日 - chendongxu - 创建。
 */
package
        com.cdx.cdxaiagent.advisor;

import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

import java.util.HashMap;
import java.util.Map;

import org.springframework.ai.chat.client.advisor.api.AdvisedRequest;
import org.springframework.ai.chat.client.advisor.api.AdvisedResponse;
import org.springframework.ai.chat.client.advisor.api.CallAroundAdvisor;
import org.springframework.ai.chat.client.advisor.api.CallAroundAdvisorChain;
import org.springframework.ai.chat.client.advisor.api.StreamAroundAdvisor;
import org.springframework.ai.chat.client.advisor.api.StreamAroundAdvisorChain;
import org.springframework.ai.chat.prompt.PromptTemplate;

/**
 * @author chendongxu
 * @since 1.0
 */
@Slf4j
public class RerAdvisor implements CallAroundAdvisor, StreamAroundAdvisor {
    private AdvisedRequest before(AdvisedRequest advisedRequest){
        // 创建包含变量的模板
        String template = """
        {re2_input_query}
        Read the question again: {re2_input_query}
        """;

        // 使用 PromptTemplate 处理模板
        PromptTemplate promptTemplate = new PromptTemplate(template);
        Map<String, Object> variables = new HashMap<>();
        variables.put("re2_input_query", advisedRequest.userText());

        String renderedText = promptTemplate.render(variables);

        return AdvisedRequest.from(advisedRequest)
                .userText(renderedText) // 使用渲染后的文本
                .userParams(variables)
                .build();
    }


    @Override
    public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {
        return chain.nextAroundCall(this.before(advisedRequest));
    }

    @Override
    public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
        return Mono.just(advisedRequest)
                .publishOn(Schedulers.boundedElastic())
                .map(request -> {
                    // 请求前处理逻辑
                    return modifyRequest(request);
                })
                .flatMapMany(request -> chain.nextAroundStream(request))
                .map(response -> {
                    // 响应处理逻辑
                    return modifyResponse(response);
                });
    }

    private AdvisedResponse modifyResponse(AdvisedResponse response) {
        return response;
    }

    private AdvisedRequest modifyRequest(AdvisedRequest request) {
        return request;
    }

    @Override
    public int getOrder() {
        return 0;
    }

    @Override
    public String getName() {
        return this.getClass().getSimpleName();
    }
}
